import cv2
import numpy as np
from skimage.feature import graycomatrix, graycoprops

class FeatureExtractor:
    def __init__(self):
        pass
    
    def extract_grid_features(self, image, grid_size):
        """
        提取图像中每个网格的特征
        
        Args:
            image: 输入图像 (RGB格式)
            grid_size: 网格大小
            
        Returns:
            网格特征字典，键为(行,列)，值为特征向量
        """
        h, w = image.shape[:2]
        features = {}
        
        for y in range(0, h, grid_size):
            for x in range(0, w, grid_size):
                # 确保网格不超出图像边界
                end_y = min(y + grid_size, h)
                end_x = min(x + grid_size, w)
                
                # 提取网格
                grid = image[y:end_y, x:end_x]
                
                # 如果网格太小，跳过
                if grid.shape[0] < 5 or grid.shape[1] < 5:
                    continue
                
                # 提取特征
                grid_features = self.extract_features(grid)
                
                # 存储特征
                features[(y // grid_size, x // grid_size)] = grid_features
        
        return features
    
    def extract_features(self, grid):
        """
        从单个网格中提取特征
        
        Args:
            grid: 网格图像
            
        Returns:
            特征向量
        """
        # 1. 颜色特征 (HSV均值)
        hsv_features = self.extract_color_features(grid)
        
        # 2. 纹理特征 (灰度共生矩阵)
        texture_features = self.extract_texture_features(grid)
        
        # 3. 形状特征 (边缘密度)
        shape_features = self.extract_shape_features(grid)
        
        # 合并所有特征
        features = np.concatenate([hsv_features, texture_features, shape_features])
        
        return features
    
    def extract_color_features(self, grid):
        """
        提取颜色特征 (HSV颜色空间)
        
        Args:
            grid: 网格图像 (RGB格式)
            
        Returns:
            HSV颜色特征向量
        """
        # 转换为HSV颜色空间
        hsv = cv2.cvtColor(grid, cv2.COLOR_RGB2HSV)
        
        # 计算每个通道的均值和标准差
        h_mean, h_std = np.mean(hsv[:,:,0]), np.std(hsv[:,:,0])
        s_mean, s_std = np.mean(hsv[:,:,1]), np.std(hsv[:,:,1])
        v_mean, v_std = np.mean(hsv[:,:,2]), np.std(hsv[:,:,2])
        
        # 计算颜色直方图
        h_hist = cv2.calcHist([hsv], [0], None, [8], [0, 180])
        s_hist = cv2.calcHist([hsv], [1], None, [4], [0, 256])
        v_hist = cv2.calcHist([hsv], [2], None, [4], [0, 256])
        
        # 归一化直方图
        h_hist = cv2.normalize(h_hist, h_hist).flatten()
        s_hist = cv2.normalize(s_hist, s_hist).flatten()
        v_hist = cv2.normalize(v_hist, v_hist).flatten()
        
        # 合并特征
        color_features = np.array([h_mean, h_std, s_mean, s_std, v_mean, v_std])
        hist_features = np.concatenate([h_hist, s_hist, v_hist])
        
        return np.concatenate([color_features, hist_features])
    
    def extract_texture_features(self, grid):
        """
        提取纹理特征 (灰度共生矩阵)
        
        Args:
            grid: 网格图像
            
        Returns:
            纹理特征向量
        """
        # 转换为灰度图
        gray = cv2.cvtColor(grid, cv2.COLOR_RGB2GRAY)
        
        # 量化灰度级别以减少计算量
        gray = (gray // 32).astype(np.uint8)
        
        # 计算灰度共生矩阵
        distances = [1]
        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]
        glcm = graycomatrix(gray, distances, angles, 8, symmetric=True, normed=True)
        
        # 提取GLCM属性
        contrast = graycoprops(glcm, 'contrast').flatten()
        dissimilarity = graycoprops(glcm, 'dissimilarity').flatten()
        homogeneity = graycoprops(glcm, 'homogeneity').flatten()
        energy = graycoprops(glcm, 'energy').flatten()
        correlation = graycoprops(glcm, 'correlation').flatten()
        
        # 合并特征
        texture_features = np.concatenate([contrast, dissimilarity, homogeneity, energy, correlation])
        
        return texture_features
    
    def extract_shape_features(self, grid):
        """
        提取形状特征 (边缘密度和方向)
        
        Args:
            grid: 网格图像
            
        Returns:
            形状特征向量
        """
        # 转换为灰度图
        gray = cv2.cvtColor(grid, cv2.COLOR_RGB2GRAY)
        
        # 计算边缘
        edges = cv2.Canny(gray, 50, 150)
        
        # 计算边缘密度
        edge_density = np.sum(edges > 0) / (edges.shape[0] * edges.shape[1])
        
        # 计算边缘方向直方图
        # 使用Sobel算子计算梯度
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        
        # 计算梯度方向
        angles = np.arctan2(sobely, sobelx) * 180 / np.pi
        
        # 创建方向直方图 (0-180度，分成6个bin)
        angle_bins = np.histogram(angles[edges > 0], bins=6, range=(0, 180))[0]
        
        # 归一化直方图
        if np.sum(angle_bins) > 0:
            angle_bins = angle_bins / np.sum(angle_bins)
        
        # 合并特征
        shape_features = np.concatenate([[edge_density], angle_bins])
        
        return shape_features
    
    def visualize_features(self, image, features, grid_size):
        """
        可视化特征
        
        Args:
            image: 原始图像
            features: 网格特征字典
            grid_size: 网格大小
            
        Returns:
            特征可视化图像
        """
        h, w = image.shape[:2]
        visualization = np.zeros((h, w, 3), dtype=np.uint8)
        
        # 为每个网格创建颜色编码
        for (row, col), feature in features.items():
            # 使用特征的前三个值来确定颜色
            # 这里简化处理，实际应用中可以使用更复杂的可视化方法
            y = row * grid_size
            x = col * grid_size
            end_y = min(y + grid_size, h)
            end_x = min(x + grid_size, w)
            
            # 使用HSV颜色特征来可视化
            h_val = int(feature[0] / 180 * 255)  # H值
            s_val = int(feature[2])              # S值
            v_val = int(feature[4])              # V值
            
            # 填充网格
            visualization[y:end_y, x:end_x] = [h_val, s_val, v_val]
        
        # 添加网格线
        for x in range(0, w, grid_size):
            cv2.line(visualization, (x, 0), (x, h), (255, 255, 255), 1)
        
        for y in range(0, h, grid_size):
            cv2.line(visualization, (0, y), (w, y), (255, 255, 255), 1)
        
        return visualization